#/usr/bin/python
# -*- coding: utf-8 -*-

## OP4 file format:
# Matrix/vector input/output files for Nastran.

from scipy import float64
import scipy.sparse as sp
import re
import time

OP4_HPATTERN = re.compile(5*r"(........)" + r"(.*)")
OP4_COLPATTERN = re.compile(3*r"(........)"+r"([\r\n])")
OP4_TESTROWPATTERN = re.compile(r"(\s*?-?\d+\.\d*E.\d*){1,5}")
OP4_ROWPATTERN = re.compile(r"(\s*?-?\d+\.\d*E.\d*)")


OP4_HFORMAT = "{COL:>8}{ROW:>8}       6       2{NAME:<8}1P,5E16.9\r\n"
OP4_COLFORMAT = "{COL:>8}{ROW:>8}{NROWS:>8}\r\n"
OP4_ROWFORMAT = ""

FLOAT = "float64"

def readOp4(op4_file,output=sp.lil_matrix):
    
    operators = {}
    
    timeref = time.time()
    while 1:
        try:
            f = open(op4_file,'r')
            break
        except IOError:
            delay = time.time()-timeref 
            if delay>5:
                raise IOError, 'Not able to open file '+op4_file

    
    line=f.readline()
    
    name = None
    col = 0
    row = 0
    nrows_in_col = 0
    
    while line != '':    
        if re.match(OP4_COLPATTERN, line): # New column
            
#            print "column header"
            inputs = re.match(OP4_COLPATTERN, line).groups()
            col = eval(inputs[0])-1
            row = eval(inputs[1])-1
            nrows_in_col = eval(inputs[2])
            
        elif re.match(OP4_TESTROWPATTERN, line): # Next row data
            if nrows_in_col == 0:
                print "Next column has not been initialized" 
               
#            print "rows in column"
            inputs = re.split(OP4_ROWPATTERN, line)
            
            
            for input in inputs:
                if input not in ['','\r\n','\n','\r']:
                    try:
                        operators[name][row,col]=float64(input)
                        row += 1
                        nrows_in_col -= 1
                    except KeyError:
                        raise KeyError, "Operator " + str(name) + " has not been initialized before."
                    except IndexError:
                        pass
                    
        elif re.match(OP4_HPATTERN, line): # New operator header
            try:
                operators[name] = output(operators[name])
            except KeyError:
                pass
            
            inputs = re.match(OP4_HPATTERN, line).groups()
#            print inputs
            name = inputs[4].strip()
            M = eval(inputs[0])
            N = eval(inputs[1])
            
            operators[name] = sp.lil_matrix((N,M),dtype=FLOAT)
            
            col = 0
            row = 0
            nrows_in_col = 0  
        else:
            break  
        line=f.readline()
    
    f.close()
            
    return operators
    
def writeOp4(op4_file,operators):
    
    try:
        f=open(op4_file,'w')
    except IOError:
        raise IOError, "Unable to open file", op4_file
    
    for op in operators.keys():
        
        name = op
        m,n = operators[op].get_shape()
        
        args = {"COL":n,"ROW":m,"NAME":name}
        # write operator's head:
        f.write(OP4_HFORMAT.format(**args))
        
        for j in xrange(n):
            
            # write new column:
            column = operators[op][:,j]
            rindex=column.nonzero()[0]
            
            try:
                args = {"COL":j+1,"ROW":rindex[0]+1,"NROWS":rindex[-1]-rindex[0]+1}
            except IndexError:
                args = False
            
            
            # nonzero column:
            if args: 
#                print "new column"
                f.write(OP4_COLFORMAT.format(**args))
                # write new row:
                cpt = 0
                rdata = []
                OP4_ROWFORMAT = ""
                specifier = {"align":">","sign":" ","width":"16","type":"E","precision":".9"}
                for i in xrange(args["NROWS"]):
                    cpt+=1
                    row = rindex[0]+i
#                    print 'row ', row, ' col ',j
                    rdata.append(operators[op][row,j])
                    OP4_ROWFORMAT = OP4_ROWFORMAT + "{0["+str(i)+"]:{align}{sign}{width}{precision}{type}}"
                    if cpt == 5:
                        cpt = 0
#                        print "writing row line..."
                        OP4_ROWFORMAT = OP4_ROWFORMAT + "\r\n" 
                        f.write(OP4_ROWFORMAT.format(rdata,**specifier))
                        OP4_ROWFORMAT = ""
                if cpt > 0:
#                    print "writing row line..."
                    OP4_ROWFORMAT = OP4_ROWFORMAT + "\r\n"
                    f.write(OP4_ROWFORMAT.format(rdata,**specifier))
                    rdata = []
                    OP4_ROWFORMAT = ""
    
    f.close()
    
    return 0

if __name__ == "__main__":
#    import os
#    testDir = "/test/"
#    relDir = os.path.relpath(testDir,os.getcwd())
#    print relDir
    
#    file = "../../test/koo1-bis.op4"
#    operators_bis = readOp4(file)

    file = "../../test/koo1.op4"
#    
    operators = readOp4(file)
    
    file = "../../test/koo1-bis.op4"
    
    writeOp4(file,operators)
#    
    operators_bis = readOp4(file)
    
    diff = operators["KOO1"]-operators_bis["KOO1"]
    
    print diff.nonzero()
    
    